package edu.protbioinfo;

import org.apache.commons.lang.StringUtils;

import edu.protbioinfo.SequenceAligner.ScoreMatrix.CellValue;
import edu.protbioinfo.SequenceAligner.ScoreMatrix.SourceDirection;
import edu.protbioinfo.Utils.Pair;




public class LinearGapGlobalSequenceAligner extends AbstractSequenceAligner {
	private ScoreMatrix scoreMatrix;
	private double gapPenality;

	public double getGapPenality() {
		return gapPenality;
	}

	public void setGapPenality(double gapPenality) {
		this.gapPenality = gapPenality;
	}

	public ScoreMatrix getScoreMatrix() {
		return scoreMatrix;
	}

	public SequenceAlignment align(final Pair<String, String> source) {
		SequenceAlignment alignment = new SequenceAlignment();
	
		if(source != null && 
		   StringUtils.isNotEmpty(source.first()) && 
		   StringUtils.isNotEmpty(source.second())){
			
//			char[] rowSeqChars = ArrayUtils.addAll(new char[]{' '}, source.second().toCharArray());
//			char[] colSeqChars = ArrayUtils.addAll(new char[]{' '}, source.first().toCharArray());
			
			scoreMatrix = new ScoreMatrix(source);
			scoreMatrix.setName("M");
			
			double maxScore = findScore(scoreMatrix.getRows() - 1, scoreMatrix.getCols() - 1);	

			alignment.setScore(maxScore);				
			traceback(alignment);
		}		
		return alignment;
	}


	protected double findScore(int i, int j) {
		double maxScore = NEGATIVE_INFINITY;
		
		Direction curDirection = null;
		
		if(i == 0 && j == 0){
			maxScore = 0.;
			curDirection = Direction.Stop;
		}
		else if(i == 0){
			maxScore = -1 * j* gapPenality;
			curDirection = Direction.Left;
		}
		else if(j == 0){
			maxScore = -1 * i * gapPenality;
			curDirection = Direction.Up;
		}
		else{
			double curScore = findScore(i - 1, j) - gapPenality;
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				curDirection = Direction.Up;
			}
			
			curScore = findScore(i, j - 1) - gapPenality;
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				curDirection = Direction.Left;
			}
			
			curScore = findScore(i - 1, j - 1) + substitutionMatrix.getValue(scoreMatrix.getRowChar(i), scoreMatrix.getColChar(j));
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				curDirection = Direction.Diagonal;
			}

		}

		scoreMatrix.setValue(i, j, new CellValue(new SourceDirection(scoreMatrix, curDirection), maxScore));		
		return maxScore;
	}

	protected void traceback(SequenceAlignment alignment) {
		int i = scoreMatrix.getRows() - 1;
		int j = scoreMatrix.getCols() - 1;
		
		Direction direction = scoreMatrix.getValue(i, j).getSourceDirection().getDirection();
	
		String firstSeq = "";
		String secondSeq = "";
	
		while(!Direction.Stop.equals(direction)){
			switch(direction){
			case Diagonal:
				firstSeq  = scoreMatrix.getColChar(j) + firstSeq;
				secondSeq = scoreMatrix.getRowChar(i) + secondSeq;
	
				i--;
				j--;
				break;
			case Up:
				firstSeq  = SequenceAlignment.BLANK + firstSeq;
				secondSeq = scoreMatrix.getRowChar(i) + secondSeq;
				
				i--;
				break;
			case Left:
				firstSeq  = scoreMatrix.getColChar(j) + firstSeq;
				secondSeq = SequenceAlignment.BLANK + secondSeq;
				
				j--;										
				break;
			}
			direction = scoreMatrix.getValue(i, j).getSourceDirection().getDirection();
		}
		
		alignment.setAlignment(new Pair<String, String>(firstSeq, secondSeq));
	}

	public static void main(String[] args){
		LinearGapGlobalSequenceAligner aligner = new LinearGapGlobalSequenceAligner();
		aligner.setGapPenality(2);
		
		aligner.setSubstitutionMatrix(new LabeledMatrix<Double>(new char[]{'A', 'C', 'D', 'E'}, new char[]{'A', 'C', 'D', 'E'}){
			{
				this.setData(new Double[]{
						2., -2., -2., -1.,
						-2., 1., 0., 0.,
						-2., 0., 2., -2.,
						-1., 0., -2., 2.
				});
			}
			
		});
		
		SequenceAlignment seqAlign = aligner.align(new Pair<String, String>("AEECA", "CDAA"));
				
		System.out.println("Score Matrix:");
		System.out.println(aligner.getScoreMatrix().toPrettyString());

		System.out.println("Optimal Alignment:");
		System.out.println(seqAlign.toPrettyString());		
	}
}
